# coding=utf-8
import matplotlib.pyplot as plt

decision_node = dict(boxstyle="sawtooth", fc="0.8")
leaf_node = dict(boxstyle="round4", fc="0.8")
arrows_args = dict(arrowstyle="<-")


def plot_node(node_txt, center_point, parent_point, node_type):
    create_plot.ax1.annotate(node_txt, xy=parent_point, xycoords='axes fraction',
                             xytext=center_point, textcoords='axes fraction',
                             va="center", ha="center", bbox=node_type,
                             arrowprops=arrows_args)


def create_plot(tree):
    fig = plt.figure(1, facecolor='white')
    fig.clf()
    axprops = dict(xticks=[], yticks=[])
    create_plot.ax1 = plt.subplot(111, frameon=False, **axprops)
    plot_ree.totalW = float(leaf_count(tree))
    plot_tree.totalD = float(tree_depth(tree))
    plot_tree.xOff = -0.5 / plot_tree.totalW
    plot_tree.yOff = 1.0
    plot_tree(tree, (0.5, 1.0), '')
    plt.show()


# count leaf of a decision tree.
def leaf_count(tree):
    leaf_num = 0
    first_str = tree.keys()[0]
    second_dict = tree[first_str]

    for key in second_dict.keys():
        if type(second_dict[key]).__name__ == 'dict':
            leaf_num += leaf_count(second_dict[key])
        else:
            leaf_num += 1

    return leaf_num


# count depth of a decision tree.
def tree_depth(tree):
    max_depth = 0
    first_str = tree.keys()[0]
    second_dic = tree[first_str]

    for key in second_dic.keys():
        if type(second_dic[key]).__name__ == 'dict':
            depth = 1 + tree_depth(second_dic[key])
        else:
            depth = 1

        if depth > max_depth:
            max_depth = depth

    return max_depth


def get_tree(i):
    trees = [{'no surfacing': {0: 'no', 1: {'flippers': {0: 'no', 1: 'yes'}}}},
             {'no surfacing': {0: 'no', 1: {'flippers': {0: {'head': {0: 'no', 1: 'yes'}}, 1: 'no'}}}}]
    return trees[i]


def plot_mid_text(centre, parent, text):
    x_mid = (parent[0] - centre[0]) / 2.0 + centre[0]
    y_mid = (parent[1] - centre[1]) / 2.0 + centre[1]
    create_plot.ax1.text(x_mid, y_mid, text)


def plot_tree(tree, parent, node_text):
    leafs_num = leaf_count(tree)
    depth = tree_depth(tree)
    first_str = tree.keys()[0]
    cntr_point = (plot_tree.xOff + (1.0 + float(leafs_num)) / 2.0 / plot_tree.totalW, plot_tree.yOff)
    plot_mid_text(cntr_point, parent, node_text)
    plot_node(first_str, cntr_point, parent, decision_node)

    second_dic = tree[first_str]
    plot_tree.yOff = plot_tree.yOff - 1.0 / plot_tree.totalD
    for key in second_dic.keys():
        if type(second_dic[key]).__name__ == 'dict':
            plot_tree(second_dic[key], cntr_point, str(key))
        else:
            plot_tree.xOff = plot_tree.xOff + 1.0 / plot_tree.totalW
            plot_node(second_dic[key], (plot_tree.xOff, plot_tree.yOff), cntr_point, leaf_node)
            plot_mid_text((plot_tree.xOff, plot_tree.yOff), cntr_point, str(key))

    plot_tree.yOff = plot_tree.yOff + 1.0 / plot_tree.totalD
